# Step 1 Import Modules
import cv2
import mediapipe as mp
import numpy as np
from matplotlib.pyplot import annotate
from mediapipe.framework.formats import landmark_pb2

mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles
myPose = mp.solutions.pose

BaseOptions = mp.tasks.BaseOptions
PoseLandmarker = mp.tasks.vision.PoseLandmarker
PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
VisionRunningMode = mp.tasks.vision.RunningMode

model_path = "models/pose_landmarker_full.task"
options = PoseLandmarkerOptions(
    base_options=BaseOptions(model_asset_path=model_path),
    running_mode=VisionRunningMode.IMAGE)

# Step 2 Create PoseLandmarker object
with PoseLandmarker.create_from_options(options) as landmarker:
    # The landmarker is initialized. Use it here.
    # ...
    # Step 3 Load input image
    # Load the input image from an image file.
    # mp_image = mp.Image.create_from_file('mn2.jpg')
    frame = cv2.imread('imgs/lyf.jpg')
    # #Load the input image from a numpy array.
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)

    # Step 4 Detect pose landmarks from the input image
    pose_landmarker_result = landmarker.detect(mp_image)

    #Step 5: Draw pos landmarks on the input image
    landmarkers = pose_landmarker_result.pose_landmarks
    pose_landmarks = landmark_pb2.NormalizedLandmarkList()
    pose_landmarks.landmark.extend([
        landmark_pb2.NormalizedLandmark(x=landmark.x, y = landmark.y, z=landmark.z, visibility=landmark.visibility)
        for landmark in landmarkers[0]
    ])
    image =mp_image.numpy_view()
    annotated_image = np.copy(image)
    mp_drawing.draw_landmarks(annotated_image, pose_landmarks, myPose.POSE_CONNECTIONS,
                              landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style())

    cv2.imshow("Image", annotated_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()








